import cv2, os, sys, time
from queue import Queue

parent_path = os.path.abspath(os.path.join(__file__, *(['..'] * 4)))
sys.path.insert(0, parent_path)

import socket, traceback
import numpy as np
from core.utils.data_io.video_reader import VideoReader
from core.utils.data_io.rtmp_stream import RtmpStream
from threading import Thread
from queue import Queue
import json
from core.utils.logger import logging
from core.utils.PIL_draw import pic_text

class TCPServer():
    def __init__(self, ip, port, buffer_size = 2) -> None:
        self.ip = ip
        self.port = port
        self.buffer_size = buffer_size
        self.server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.server_socket.bind((ip,port))
        self.server_socket.listen()
        self.buffer = Queue(maxsize=buffer_size)
        self.last_data = {}
        self.max_retry = 5
        self.status = 1 # 0:done 1:standby 2: running
        
    def _running(self,):
        self.status = 2
        try_num = 0
        while self.status:
            client_socket, clientAddr = self.server_socket.accept()
            while self.status:
                # 接收对方发送过来的数据
                try:
                    recv_data = client_socket.recv(2048*5)
                except:
                    logging("|接受收据失败, 关闭TCP链接|")
                    client_socket.close()
                    break

                if len(recv_data)<=0: # 空数据关闭链接
                    time.sleep(1)
                    try_num+=1
                    if try_num>=self.max_retry:
                        client_socket.close()
                        try_num = 0
                        logging.error("|空数据次数到达最大值, 关闭TCP链接|")
                        client_socket.close()
                        break
                recv_data = recv_data.decode()
                try:
                    js_data = json.loads(recv_data)
                except:
                    logging.error("json load 失败")
                    continue
                try:
                    self.buffer.put_nowait(js_data)
                except:
                    pass

    def start(self):
        thread = Thread(target = self._running)
        thread.start()
    
    def get_data(self):
        try:
            self.last_data = self.buffer.get_nowait()
        except:
            pass
        finally:
            return self.last_data
        
class TcpSteamPush():
    def __init__(self, ip, port) -> None:
        self.status = 0
        self.server = TCPServer(ip,port)
        self.reader = VideoReader("real_time")
        self.server.start()

    def rtmp_init(self, pull_url, push_url):
        self.reader.init_cap(pull_url)
        self.writer = RtmpStream(push_url,1280,720,15)
        self.writer.start()
        self.reader.start()
        for _ in range(5):
            time.sleep(2)
            ret,_ = self.reader.read()
            if ret:
                break
        assert ret == True
        self.status = 1
        
    
    def run_loop(self):
        try:
            wait_time = 1/15
            while self.status == 2:
                runing_start = time.time()
                ret, frame = self.reader.read()
                if not ret:
                    continue
                frame = cv2.resize(frame,(1280,720))
                bboxes = self.server.get_data()
                # for key, det in bboxes.items():
                    # cls, score, x1,y1,x2,y2 = bbox
                    # frame = cv2.rectangle(frame,[int(x1),int(y1)],[int(x2),int(y2)],[0,0,255],2)
                if bboxes.get("data",False):
                    results = bboxes["data"]
                    points = []
                    for result in results: # 每张图
                        if result is None:
                            continue
                        for object in result: # 每个目标
                            points.append(object[0])
                            point = object[0][3]
                            point = [int(x) for x in point]
                            point[1]+=30
                            txt = object[1][0]
                            frame = pic_text(frame, txt, point, (0,0,255))
                            # score = object[1][1]
                    points = np.array(points,dtype=np.int64).reshape(-1,4,2)
                    cv2.polylines(frame, points, True, (0,255,0), 2)
                self.writer.put(frame)
                time.sleep(max(wait_time-(time.time()-runing_start),0))
        except:
            logging.error(traceback.format_exc())
        finally:
            logging.info("|TcpSteamPush run_loop done.|")

    def start(self):
        thread = Thread(target=self.run_loop)
        self.status = 2
        thread.start()

    def close(self):
        self.writer.close()
        self.reader.close()
        self.status = 2

if __name__ == "__main__":
    import cv2, time
    wait_time = 1/20
    reader = VideoReader("real_time")
    pull_url = "rtmp://218.94.155.86:1936/live/livestream_74_1"
    push_url = "rtmp://218.94.155.86:1936/live/livestream_74_3"
    ip = "0.0.0.0"
    port = 2692
    
    tcp_steam_pusher = TcpSteamPush(ip,port)
    tcp_steam_pusher.rtmp_init(pull_url,push_url)
    tcp_steam_pusher.start()
    while 1:
        time.sleep(1)